package org.zetaframework.satoken.interceptor;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.web.servlet.HandlerInterceptor;
import org.zetaframework.core.constants.Constants;
import org.zetaframework.core.utils.MdcContextUtil;

/**
 * 链路追踪拦截器
 *
 * @author gcc
 */
public class TraceInterceptor implements HandlerInterceptor {
    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    /**
     * 每次请求之前触发的方法
     *
     * @param request HttpServletRequest
     * @param response HttpServletResponse
     * @param handler Object
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        // 从请求头中获取traceId
        String traceId = request.getHeader(Constants.TRACE_ID);

        // 如果为空，生成新的traceId
        if (StrUtil.isBlank(traceId)) {
            traceId = IdUtil.fastSimpleUUID();
            logger.warn("从请求头中获取到的traceId为空，生成新traceId：{}", traceId);
        }

        // 设置到MDC中，用于日志追踪
        MDC.put(Constants.TRACE_ID, traceId);
        // 设置到ThreadLocal中，用于异步环境获取traceId
        MdcContextUtil.setTraceId(traceId);
        return true;
    }

    /**
     * 执行完Controller之后，要做的事
     *
     * @param request HttpServletRequest
     * @param response HttpServletResponse
     * @param handler Object
     * @param ex Exception
     */
    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        // 清理MDC中的traceId，防止内存泄漏
        MDC.clear();
        // 清空ThreadLocal的值，防止下次请求时获取到的值是旧数据，同时也能防止内存溢出
        MdcContextUtil.remove();
    }
}
